from traindatawithmysql import aidatatemplateCount as ACount
from traindatawithmysql import aidatatemplateids as GCountTIds
from traindatawithmysql import aidataassignids as GDataWithTIds
from traindatawithmysql import GetAIDataWithAssignIds
from FunctionTool import RecursionLoop as Loop
from traindata import TrainDataColumns

from concurrent.futures import ProcessPoolExecutor as PPool

import pandas as pd
import json
import os, math, threading, time, gc

const_groupId = 'template_id'
const_space = "space"
const_input = "input"
const_item = "item"
const_attrs = "attrs"
const_weight = "weight"
const_target = "target"
config = {
    const_space: ["space_name", "template_group"],
    const_item: ["gmodel_cat_name"],
    const_attrs: ["tag_2", "tag_3", "tag_4", "tag_6",
                  "tag_7", "tag_8", "tag_10", "tag_11",
                  "tag_12", "tag_13", "tag_14", "tag_15",
                  "tag_16", "tag_17", "tag_18", "tag_31",
                  "tag_43", "tag_44", "tag_45", "tag_49",
                  "tag_54", "tag_55", "tag_56", "tag_57",
                  "tag_58", "tag_59", "tag_60", "tag_62",
                  "tag_63", "tag_64", "tag_65", "color_hue",
                  "color_lightness", "color_purity", "gmodel_id_tag",
                  "decorate_level", "base_color_id", "colour_scheme", "base_color_show"],
    const_weight: ["weight"]
}
try:
    config[const_attrs] = TrainDataColumns()
except:
    pass


#
def AutoTrain():
    totalCount = ACount()
    everyCount = 1000
    whileCount = math.ceil(totalCount / everyCount)

    dic, extra = {}, {}
    with PPool(max_workers=3) as pool:
        tasks = [pool.submit(Train, i * everyCount, min(everyCount, totalCount - i * everyCount)) for i in
                 range(whileCount)]
        pool.shutdown(True)
        time.sleep(1)
        gc.collect()
        print("处理结果")
        begintime = time.time()
        for task in tasks:
            _dic, _extra = task.result()
            MergeDict(dic, _dic)
            MergeExtra(extra, _extra)
        print("结果处理完成", time.time() - begintime)
    print("转换数据")
    time.sleep(1)
    gc.collect()
    begintime = time.time()
    dic = ParseKey5ToKey4(dic)
    print("转换数据完成", time.time() - begintime)
    return dic, extra[const_space], extra[const_target], extra[const_attrs]


def Train(index, count):
    global const_groupId, config, const_space, const_item, const_attrs, const_weight
    setIds = GCountTIds(index, count)
    dic, extra = {}, {}

    def _action(query):
        nonlocal dic, extra, setIds
        i = 0
        for id in setIds:
            data = query(set([str(id)]))
            vail = filterVailColumn(data)
            result = {}
            trainWithVailDict(vail, result)
            MergeResultToDict(result, dic, extra)
            print(index + i, data.index.size)
            i += 1

    GetAIDataWithAssignIds(_action)
    return dic, extra


def trainWithVailDict(vail, result):
    for i, val in vail.items():
        TrainWithRow(val, result, i)


def filterVailColumn(data):
    global config, const_attrs, const_space, const_item, const_weight
    attrs = config[const_attrs] + config[const_space] + config[const_item] + config[const_weight]
    vail = {}
    for index, row in data.iterrows():
        dic = {}
        vail[row.get("id", index)] = dic
        for attr in attrs:
            if attr not in row:
                continue
            try:
                val = row.get(attr)
                val = HandleSpecialKey(attr, val)
                if val is None:
                    val = ""
                val = str(val)
                if val.find("b'") != -1:
                    val = val.replace("b'", "")
                    val = val.replace("'", "")
            except:
                val = ""
            if val == "":
                continue
            dic[attr] = val
    return vail


# 递归循环
def RecursionLoop(array, handle):
    if handle is None:
        return

    cnt = len(array)

    def _recursion(index):
        nonlocal array, cnt, handle
        if index >= cnt:
            return

        temp = array[index]

        for t in temp:
            if handle(t, index):
                _recursion(index + 1)

    _recursion(0)


def HandleSpecialKey(key, val):
    try:
        if key == "base_color_show" and val and val != "":
            val = val.split("^")[0]
    except Exception as e:
        print(e)
    return val


# 处理一条数据
def TrainWithRow(row, result, id):
    global config, const_space, const_item, const_attrs, const_weight
    array = [config[const_space], config[const_item], config[const_attrs], config[const_weight]]
    cnt = len(array)

    rets = []

    def _handle(key, index):
        nonlocal row, result, cnt
        val = ""
        try:
            val = row.get(key)
            val = HandleSpecialKey(key, val)
            if val is None:
                val = ""
            val = str(val)
            if val.find("b'") != -1:
                val = val.replace("b'", "")
                val = val.replace("'", "")
        except:
            val = ""
        retVal = (key, val)
        if index == len(rets):
            rets.append(retVal)
        else:
            rets[index] = retVal

        if index + 1 == cnt and val != "":
            MergeRetToResult(rets, result, id)

        return val != ""

    RecursionLoop(array, _handle)


# 结果合并到result
def MergeRetToResult(rets, result, id):
    id = str(id)
    tempResult = result
    weight = 1
    try:
        weight = int(rets[3])
    except:
        weight = 1

    if rets[0] not in tempResult:
        tempResult[rets[0]] = {}
    tempResult = tempResult[rets[0]]
    if rets[1] not in tempResult:
        tempResult[rets[1]] = {}
    tempResult = tempResult[rets[1]]

    if rets[2] not in tempResult:
        tempResult[rets[2]] = [0, set()]
    tempResult[rets[2]][0] += weight
    tempResult[rets[2]][1].add(id)


# 合并成一条数据
def MergeResultToDict(result, dict, extra):
    global const_space, const_attrs, const_target
    if result is None:
        return
    for key, val in result.items():
        for k, v in val.items():
            for _k, _v in val.items():
                if k == _k:
                    continue
                for vk, vv in v.items():
                    for _vk, _vv in _v.items():
                        # 正常
                        fkey, fweight, ftitle = (key[1], k[1], vk[1], _k[1], _vk[1]), vv[0] * _vv[0], (
                        key[0], k[0], vk[0], _k[0], _vk[0])
                        lefId = set() | (vv[1])
                        rigId = set() | (_vv[1])
                        if fkey in dict:
                            fweight += dict[fkey][0]
                            lefId.update(dict[fkey][6])
                            rigId.update(dict[fkey][7])
                        dict[fkey] = (fweight,) + ftitle + (lefId, rigId)
                        if const_space not in extra:
                            extra[const_space] = set()
                        if const_target not in extra:
                            extra[const_target] = set()
                        if const_attrs not in extra:
                            extra[const_attrs] = set()
                        extra[const_space].add(key[1])
                        extra[const_target].add(k[1])
                        extra[const_target].add(_k[1])
                        extra[const_attrs].add(vk[1])
                        extra[const_attrs].add(_vk[1])

                        # # input 缺省
                        # fkey, fweight, ftitle = (key[1], "*", vk[1], _k[1], _vk[1]), vv * _vv, (key[0], "*", vk[0], _k[0], _vk[0])
                        # if fkey in dict:
                        #     fweight += dict[fkey][0]
                        # dict[fkey] = (fweight,) + ftitle
                        # # space 缺省
                        # fkey, fweight, ftitle = ("*", k[1], vk[1], _k[1], _vk[1]), vv * _vv, ("*", k[0], vk[0], _k[0], _vk[0])
                        # if fkey in dict:
                        #     fweight += dict[fkey][0]
                        # dict[fkey] = (fweight,) + ftitle


def MergeDict(dic, _dic):
    for _key, _val in _dic.items():
        __val = list(_val)
        if _key in dic:
            __val[0] += dic[_key][0]
            __val[len(__val) - 1].update(dic[_key][len(dic[_key]) - 1])
        dic[_key] = tuple(__val)


def MergeExtra(extra, _extra):
    for _key, _val in _extra.items():
        if _key in extra:
            extra[_key] = extra[_key].union(_val)
        else:
            extra[_key] = _val


# key5 => key4
def ParseKey5ToKey4(dic):
    result = {}
    for key, val in dic.items():
        k = tuple(list(key)[0:4])
        v = (key[4],) + val
        if k not in result:
            result[k] = []
        result[k].append(v)
    return result


# 训练
if __name__ == "__main__":
    begintime = time.time()
    dic, spaces, targets, attrs = AutoTrain()
    print(time.time() - begintime)
